package de.uni_stuttgart.ipvs.ids.communicationLib;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.security.NoSuchAlgorithmException;

import de.uni_stuttgart.ipvs.ids.communicationLib.crypto.InvalidSignatureException;
import de.uni_stuttgart.ipvs.ids.communicationLib.crypto.PaddingException;
import de.uni_stuttgart.ipvs.ids.communicationLib.crypto.XorCryptoProvider;

/**
 * A wrapper around a regular DatagramSocket providing security functions.
 * 
 * @author Team 3/4 Elkhan Dadashov, Ahmed Labib, Dhaval A Shah, Patrick Stortz
 */
public class SLSDatagramSocket {

	/** The underlying DatagramSocket to use */
	DatagramSocket wrappedSocket;
	/** The XorCryptoProvider to use */
	private XorCryptoProvider cryptoProvider;

	/**
	 * Creates a new SLSDatagramSocket.
	 * @param wrappedSocket the underlying DatagramSocket to use.
	 * @param cryptoProvider the XorCryptoProvider to use.
	 */
	public SLSDatagramSocket(DatagramSocket wrappedSocket,
			XorCryptoProvider cryptoProvider) {
		this.wrappedSocket = wrappedSocket;
		this.cryptoProvider = cryptoProvider;
	}

	/**
	 * Passes the contents of the given DatagramPacket to the cryptoPRovider for
	 * decryption. The contents of the Packet are replaced with the decrypted
	 * data.
	 * 
	 * @param p
	 *            The DatagramPacket whose data is to be decrypted.
	 * @throws IOException
	 *             If an error during decryption occurs.
	 * @throws InvalidSignatureException If the signature does not match the payload.
	 */
	protected void decryptPacket(DatagramPacket p) throws IOException,
			InvalidSignatureException {

		// get the signature length from the provider
		int sigLength = this.cryptoProvider.getSignatureLength(p.getData(),
				p.getOffset(), p.getLength());

		// extract the cipher text and signature from the payload
		byte[] cipher = new byte[p.getLength() - sigLength];
		byte[] signature = new byte[sigLength];
		System.arraycopy(p.getData(), 0, cipher, 0, cipher.length);
		System.arraycopy(p.getData(), p.getLength() - sigLength, signature, 0,
				signature.length);

		// check the signature
		boolean signatureValid = false;
		try {
			signatureValid = this.cryptoProvider.checkSignature(cipher,
					signature);
		} catch (NoSuchAlgorithmException e) {
			System.err
					.println("decryptPacket() - Signature could not be generated!"); // DEBUG
			throw new IOException("Signature could not be generated!");
		} catch (PaddingException e) {
			System.err
					.println("decryptPacket() - Signature could not be generated!"); // DEBUG
			throw new IOException("Signature could not be generated!");
		}

		if (signatureValid == false) {
			System.out.println("\t> decryptPacket() - Signature is invalid! <"); // DEBUG
			throw new InvalidSignatureException();
		}

		// decrypt the cipher text
		byte[] plain = null;
		try {
			plain = this.cryptoProvider.decrypt(cipher);
		} catch (PaddingException e1) {
			System.err
					.println("decryptPacket() - Cipher text could not be decrypted!"); // DEBUG
			throw new IOException("Cipher text could not be decrypted!");
		}

		// set the plain text as payload
		p.setData(plain);

		System.out.println("\t> decryptPacket() - Success! <"); // DEBUG
	}

	/**
	 * Passes the contents of the given DatagramPacket to the cryptoPRovider for
	 * encryption. The contents of the Packet are replaced with the encryption
	 * data.
	 * 
	 * @param p
	 *            The DatagramPacket whose data is to be encryption.
	 * @throws IOException
	 *             If an error during encryption occurs.
	 */
	protected void encryptPacket(DatagramPacket p) throws IOException {

		// encrypt the plain text - E(msg)
		byte[] cipher = null;
		try {
			cipher = this.cryptoProvider.encrypt(p.getData());
		} catch (PaddingException e) {
			System.err
					.println("encryptPacket() - Plain text could not be encrypted!"); // DEBUG
			throw new IOException("Plain text could not be encrypted!");
		}

		// generate the signature - Sig(E(msg))=E(SHA1(E(msg)))
		byte[] signature = null;
		try {
			signature = this.cryptoProvider.generateSignature(cipher);
		} catch (NoSuchAlgorithmException e) {
			System.err
					.println("encryptPacket() - Signature could not be generated!"); // DEBUG
			throw new IOException("Signature could not be generated!");
		} catch (PaddingException e) {
			System.err
					.println("encryptPacket() - Signature could not be generated!"); // DEBUG
			throw new IOException("Signature could not be generated!");
		}

		// concatenate the cipher and signature and put it into the payload
		byte[] data = new byte[cipher.length + signature.length];
		System.arraycopy(cipher, 0, data, 0, cipher.length);
		System.arraycopy(signature, 0, data, cipher.length, signature.length);
		p.setData(data, 0, data.length);

		System.out.println("\t> encryptPacket() - Success! <"); // DEBUG
	}

	/**
	 * Receives a DatagramPacket and decrypts it.
	 * @param p the DatagramPacket to use for receiving.
	 * @throws IOException If receiving or decrypting fails.
	 */
	public synchronized void receive(DatagramPacket p) throws IOException {
		wrappedSocket.receive(p);
		decryptPacket(p);
	}

	/**
	 * Encrypts and sends a DatagramPacket.
	 * @param p the DatagramPacket to encrypt and send.
	 * @throws IOException If encrypting or sending fails.
	 */
	public void send(DatagramPacket p) throws IOException {
		encryptPacket(p);
		wrappedSocket.send(p);
	}

	/**
	 * Encrypts and sends a DatagramPacket. Flips a random byte within the cipher text part.
	 * @param p the DatagramPacket to encrypt, to manipulate and to send.
	 * @throws IOException If encrypting or sending fails.
	 */
	public void sendWithBrokenMessage(DatagramPacket p) throws IOException {
		encryptPacket(p);
		// manipulate the message
		int signatureLength = cryptoProvider.getSignatureLength(p.getData(),
				p.getOffset(), p.getLength());
		int byteToFlip = (int) Math.rint(Math.random()
				* (p.getLength() - signatureLength));
		p.getData()[byteToFlip] ^= 1;
		wrappedSocket.send(p);
	}

	/**
	 * Encrypts and sends a DatagramPacket. Flips a random byte within the signature part.
	 * @param p the DatagramPacket to encrypt, to manipulate and to send.
	 * @throws IOException If encrypting or sending fails.
	 */
	public void sendWithBrokenSignature(DatagramPacket p) throws IOException {
		encryptPacket(p);
		// manipulate the signature
		int signatureLength = cryptoProvider.getSignatureLength(p.getData(),
				p.getOffset(), p.getLength());
		int byteToFlip = (int) Math.rint(Math.random() * (signatureLength));
		byteToFlip += p.getLength() - signatureLength;
		p.getData()[byteToFlip] ^= 1;
		wrappedSocket.send(p);
	}

}
